import os
import numpy as np
import matplotlib.pyplot as plt
from copy import deepcopy

#save_path = os.path.dirname(__file__)
save_path = os.getcwd()
POSITION = (1, 0)
VELOCITY = (0, 5)
precision_default = 1000
T_times = 10
GM = 25

def Simulation(position, velocity, t, method, level=2):
    x, y = position
    vx, vy = velocity
    r = np.sqrt(np.square(x) + np.square(y))
    ax = -x * np.power(r, -(level + 1)) * GM
    ay = -y * np.power(r, -(level + 1)) * GM
    if method == 'EulerCromer':
        vx += ax * t
        vy += ay * t
        x += vx * t
        y += vy * t
    elif method == 'Euler':
        x += vx * t
        y += vy * t
        vx += ax * t
        vy += ay * t
    elif method == 'EulerRichardson':
        x_half = x + vx * t / 2
        y_half = y + vy * t / 2
        vx_half = vx + ax * t / 2
        vy_half = vy + ay * t / 2
        r = np.sqrt(np.square(x_half) + np.square(y_half))
        ax_half = -x_half * np.power(r, -(level + 1)) * GM
        ay_half = -y_half * np.power(r, -(level + 1)) * GM
        x += vx_half * t
        y += vy_half * t
        vx += ax_half * t
        vy += ay_half * t
    return (x, y), (vx, vy)

def draw(position_init, velocity_init, t, method, name, indexes, level=2):
    x = np.array([position_init[0]])
    y = np.array([position_init[1]])
    velocity = deepcopy(velocity_init)
    for i in indexes[:-1]:
        position, velocity = Simulation((x[-1], y[-1]), velocity, t, method, level)
        x = np.append(x, position[0])
        y = np.append(y, position[1])
    plt.subplot(211)
    plt.plot(x, y)
    plt.axis('equal')
    plt.subplot(212)
    plt.plot(indexes / T, np.sqrt(np.square(x) + np.square(y)))
    plt.grid()
    plt.savefig(os.path.join(save_path, name))
    plt.clf()

if __name__ == "__main__":
    T = 2 * np.pi * np.sqrt(np.square(POSITION[0]) + np.square(POSITION[1])) / np.sqrt(np.square(VELOCITY[0]) + np.square(VELOCITY[1]))
    t = T / precision_default
    indexes = np.linspace(0, T * T_times, precision_default * T_times)
    for method in ['EulerCromer', 'Euler', 'EulerRichardson']:
        draw(POSITION, VELOCITY, t, method, "1_" + method + ".png", indexes)
    
    for method in ['EulerCromer', 'EulerRichardson']:
        for precision in [100, 1000, 10000]:            
            t = T / precision
            indexes = np.linspace(0, T * T_times, precision * T_times)
            for level in [1.95, 2.05]:
                draw(POSITION, VELOCITY, t, method, "2_" + method + "_" + str(precision) + "_" + str(level) + ".png", indexes, level)
            for velocity in [VELOCITY, (0, 5.1), (0, 4.9)]:
                draw(POSITION, velocity, t, method, "3_" + method + "_" + str(precision) + "_" + str(velocity[1]) + ".png", indexes, 3)